"""Federated client implementation."""
import pickle
import torch
from typing import Dict, Any
from .protocols.auth import ECDSAAuth, AuthScheme
from .protocols import privacy
from . import model as model_mod
from . import dataset as ds_mod

class Client:
    def __init__(self, client_id: str, dataloader, auth_scheme: AuthScheme = None, device: str = 'cpu', privacy_layer: privacy.PrivacyLayer = None):
        self.id = client_id
        self.dataloader = dataloader
        self.auth = auth_scheme or ECDSAAuth()
        self.sk, self.pk = self.auth.keygen()
        self.device = device
        self.model = model_mod.get_model().to(device)
        self.privacy_layer = privacy_layer

    def train_local(self, epochs=1, lr=0.01):
        delta, num_samples = ds_mod.simulate_local_update(self.model, self.dataloader, epochs=epochs, lr=lr, device=self.device)
        return delta, num_samples

    def package_and_sign(self, delta_state: Dict[str, torch.Tensor]):
        """Serialize delta and sign. Returns (signature, message_bytes)."""
        message = pickle.dumps({k: v.cpu() for k,v in delta_state.items()})
        sig, _ = self.auth.sign(self.sk, message)
        return sig, message

    def get_public_key(self):
        return self.pk
